#coding=utf8
from __future__ import division
import torch
import os,time,datetime
from torch.autograd import Variable
import logging
import torch.nn.functional as F
import numpy as np
from math import ceil
import copy
from metrics import cal_mAP
from predicting import predict


def dt():
    return datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')

def trainlog(logfilepath, head='%(message)s'):
    logger = logging.getLogger('mylogger')
    logging.basicConfig(filename=logfilepath, level=logging.INFO, format=head)

    console = logging.StreamHandler()
    console.setLevel(logging.INFO)
    formatter = logging.Formatter(head)
    console.setFormatter(formatter)
    logging.getLogger('').addHandler(console)


def train(model,
          epoch_num,
          start_epoch,
          optimizer,
          criterion,
          exp_lr_scheduler,
          multi_loader,
          val_data_set,
          val_data_loader,
          save_dir,
          augloss=False,
          print_inter=200,
          val_inter=3500
          ):

    best_model_wts = model.state_dict()
    best_acc = 0.0
    best_mAP = 0.0


    multi_loader.VirEpoch = start_epoch-1
    while multi_loader.VirEpoch < epoch_num:
        exp_lr_scheduler.step(multi_loader.VirEpoch)
        model.train()

        # which loader to provide data
        i = multi_loader.Step % multi_loader.num_datasets
        batch = multi_loader.next(i)
        inputs, attr, attr_mask, labels, labels_str = batch

        inputs = Variable(inputs.cuda())
        labels = Variable(labels.cuda())

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward
        if augloss:
            outputs, aux = model(inputs)
            loss = criterion(outputs, labels) + criterion(aux, labels)
        else:
            # print inputs
            outputs = model(inputs)
            loss = criterion(outputs, labels)

        loss.backward()
        optimizer.step()

        # batch loss
        if multi_loader.Step % print_inter == 0:
            attr_mask = Variable(attr_mask.cuda())
            _, preds = torch.max(outputs * attr_mask, 1)
            _, labels = torch.max(labels, 1)

            batch_corrects = torch.sum((preds == labels)).data[0]
            batch_acc = batch_corrects / (labels.size(0))

            logging.info('%s [%d-%d] | batch-loss: %.3f | acc@1: %.3f'
                         % (dt(), multi_loader.VirEpoch, multi_loader.Step,
                            loss.data[0], batch_acc))

        # do eval
        if multi_loader.Step % val_inter == 0:
            logging.info('--' * 30)
            logging.info('current lr:%s' % exp_lr_scheduler.get_lr())
            val_mAP, val_acc = predict(model, val_data_set, val_data_loader, whileTraing=True)

            if val_mAP > best_mAP:
                best_mAP = val_mAP
                best_acc = val_acc
                best_model_wts = copy.deepcopy(model.state_dict())

            # save model
            save_path = os.path.join(save_dir,
                                     'weights-%d-%d-[%.4f]-[%.4f].pth' % (
                                         multi_loader.VirEpoch, multi_loader.Step,
                                         val_acc, val_mAP))
            torch.save(model.state_dict(), save_path)
            logging.info('saved model to %s' % (save_path))
            logging.info('--' * 30)

    # save best model
    save_path = os.path.join(save_dir,
                             'bestweights-[%.4f]-[%.4f].pth' % (best_acc,best_mAP))
    torch.save(best_model_wts, save_path)
    logging.info('saved model to %s' % (save_path))


    return best_acc,best_model_wts